module latlon_halo_mod

  use mpi
  use const_mod
  use latlon_mesh_mod

  implicit none

  private

  public latlon_halo_type

  integer, parameter :: cross_proc_halo = 1
  integer, parameter :: cross_comm_halo = 2
  integer, parameter :: inner_halo = 3
  integer, parameter :: nest_halo = 4

  type latlon_halo_type
    integer :: comm = MPI_COMM_NULL
    integer :: host_id = MPI_PROC_NULL
    integer :: proc_id = MPI_PROC_NULL
    integer :: orient = 0
    integer :: dtype = 0
    integer :: type = 0
    integer :: send_type_2d = MPI_DATATYPE_NULL
    integer :: recv_type_2d = MPI_DATATYPE_NULL
    integer :: send_type_3d = MPI_DATATYPE_NULL
    integer :: recv_type_3d = MPI_DATATYPE_NULL
  contains
    procedure :: init => latlon_halo_init
    final :: latlon_halo_final
  end type latlon_halo_type

contains

  subroutine latlon_halo_init(this, mesh, orient, dtype, host_id, ngb_proc_id, lon_ibeg, lon_iend, lat_ibeg, lat_iend)

    class(latlon_halo_type), intent(out) :: this
    type(latlon_mesh_type), intent(in) :: mesh
    integer, intent(in) :: orient
    integer, intent(in) :: dtype
    integer, intent(in) :: host_id
    integer, intent(in) :: ngb_proc_id
    integer, intent(in), optional :: lon_ibeg
    integer, intent(in), optional :: lon_iend
    integer, intent(in), optional :: lat_ibeg
    integer, intent(in), optional :: lat_iend

    integer full_lon_ibeg, full_lon_iend
    integer full_lat_ibeg, full_lat_iend
    integer array_size(3)
    integer send_subarray_size(3)
    integer recv_subarray_size(3)
    integer send_subarray_start(3)
    integer recv_subarray_start(3)
    integer num_lev
    integer i, j, k, ierr

    this%proc_id = ngb_proc_id

    this%host_id = host_id
    this%dtype = dtype

    if (present(lon_ibeg) .and. present(lon_iend)) then
      full_lon_ibeg = lon_ibeg - mesh%lon_lb
      full_lon_iend = lon_iend - mesh%lon_lb
    else if (orient == west) then
      full_lon_ibeg = 0
      full_lon_iend = mesh%lon_halo_width - 1
    else if (orient == east) then
      full_lon_ibeg = mesh%num_lon + mesh%lon_halo_width
      full_lon_iend = full_lon_ibeg + mesh%lon_halo_width - 1
    end if
    if (present(lat_ibeg) .and. present(lat_iend)) then
      full_lat_ibeg = lat_ibeg - mesh%lat_lb
      full_lat_iend = lat_iend - mesh%lat_lb
    else if (orient == south) then
      full_lat_ibeg = 0
      full_lat_iend = mesh%lat_halo_width - 1
    else if (orient == north) then
      full_lat_ibeg = mesh%num_lat + mesh%lat_halo_width
      full_lat_iend = full_lat_ibeg + mesh%lat_halo_width - 1
    end if

    ! NOTE: MPI array index starts from zero.

    !                          wx                          nx                          wx      
    !                          |                           |                           |
    !                  |---------------|---------------------------------------|---------------|
    !                  |_______________|_______________________________________|_______________|__ 
    !                  |\\\\\\\|\\\\\\\|\\\\\\\|\\\\\\\|\\\\\\\|\\\\\\\|\\\\\\\|\\\\\\\|\\\\\\\|  |
    !         wy + ny -|\\\\\\\|\\\\\\\|\\\\\\\|\\\\\\\|\\\\\\\|\\\\\\\|\\\\\\\|\\\\\\\|\\\\\\\|  |- wy
    !                  |\\\\\\\|\\\\\\\|\\\\\\\|\\\\\\\|\\\\\\\|\\\\\\\|\\\\\\\|\\\\\\\|\\\\\\\|  |
    !                  |_______|_______|_______________________________________|_______________|__|
    !                  |///////|///////|       |       |       |       |       |///////|///////|  |
    !     wy + ny - 1 -|///////|///////|       |       |       |       |       |///////|///////|  |
    !                  |///////|///////|       |       |       |       |       |///////|///////|  |
    !                  |_______|_______|_______|_______|_______|_______|_______|_______|_______|  |
    !                  |///////|///////|       |       |       |       |       |///////|///////|  |
    !                  |///////|///////|       |       |       |       |       |///////|///////|  |- ny
    !                  |///////|///////|       |       |       |       |       |///////|///////|  |
    !                  |_______|_______|_______|_______|_______|_______|_______|_______|_______|  |
    !                  |///////|///////|       |       |       |       |       |///////|///////|  |
    !              wy -|///////|///////|       |       |       |       |       |///////|///////|  |
    !                  |///////|///////|       |       |       |       |       |///////|///////|  |
    !                  |_______|_______|_______|_______|_______|_______|_______|_______|_______|__|
    !                  |\\\\\\\|\\\\\\\|\\\\\\\|\\\\\\\|\\\\\\\|\\\\\\\|\\\\\\\|\\\\\\\|\\\\\\\|  |
    !               0 -|\\\\\\\|\\\\\\\|\\\\\\\|\\\\\\\|\\\\\\\|\\\\\\\|\\\\\\\|\\\\\\\|\\\\\\\|  |- wy
    !                  |\\\\\\\|\\\\\\\|\\\\\\\|\\\\\\\|\\\\\\\|\\\\\\\|\\\\\\\|\\\\\\\|\\\\\\\|  |
    !                  |_______|_______|_______|_______|_______|_______|_______|_______|_______|__|
    !                      |               |                       |       |       |
    !                      0               wx                      | wx + nx - 1   |
    !                                                              |            wx + nx
    !                                                              nx

    this%orient = orient
    this%type = cross_proc_halo
    num_lev = 1

    array_size = [mesh%num_lon+2*mesh%lon_halo_width,mesh%num_lat+2*mesh%lat_halo_width,num_lev]
    send_subarray_size = [full_lon_iend-full_lon_ibeg+1,full_lat_iend-full_lat_ibeg+1,num_lev]
    recv_subarray_size = send_subarray_size
    select case (orient)
    case (west)
      send_subarray_start = [full_lon_iend+1,full_lat_ibeg,0]
      recv_subarray_start = [full_lon_ibeg  ,full_lat_ibeg,0]
    case (east)
      send_subarray_start = [full_lon_ibeg-mesh%lon_halo_width,full_lat_ibeg,0]
      recv_subarray_start = [full_lon_ibeg                    ,full_lat_ibeg,0]
    case (south)
      send_subarray_start = [full_lon_ibeg,full_lat_iend+1,0]
      recv_subarray_start = [full_lon_ibeg,full_lat_ibeg  ,0]
    case (north)
      send_subarray_start = [full_lon_ibeg,full_lat_ibeg-mesh%lat_halo_width,0]
      recv_subarray_start = [full_lon_ibeg,full_lat_ibeg                    ,0]
    end select
    call MPI_TYPE_CREATE_SUBARRAY(3, array_size, send_subarray_size, &
                                  send_subarray_start, MPI_ORDER_FORTRAN, dtype, &
                                  this%send_type_3d, ierr)
    call MPI_TYPE_COMMIT(this%send_type_3d, ierr)
    call MPI_TYPE_CREATE_SUBARRAY(3, array_size, recv_subarray_size, &
                                  recv_subarray_start, MPI_ORDER_FORTRAN, dtype, &
                                  this%recv_type_3d, ierr)
    call MPI_TYPE_COMMIT(this%recv_type_3d, ierr)

    call MPI_TYPE_CREATE_SUBARRAY(2, array_size(1:2), send_subarray_size(1:2), &
                                  send_subarray_start(1:2), MPI_ORDER_FORTRAN, dtype, &
                                  this%send_type_2d, ierr)
    call MPI_TYPE_COMMIT(this%send_type_2d, ierr)
    call MPI_TYPE_CREATE_SUBARRAY(2, array_size(1:2), recv_subarray_size(1:2), &
                                  recv_subarray_start(1:2), MPI_ORDER_FORTRAN, dtype, &
                                  this%recv_type_2d, ierr)
    call MPI_TYPE_COMMIT(this%recv_type_2d, ierr)

  end subroutine latlon_halo_init

  subroutine latlon_halo_final(this)

    type(latlon_halo_type), intent(inout) :: this

    integer i, j, k
    integer ierr

    if (this%send_type_3d /= MPI_DATATYPE_NULL) call MPI_TYPE_FREE(this%send_type_3d, ierr)
    if (this%recv_type_3d /= MPI_DATATYPE_NULL) call MPI_TYPE_FREE(this%recv_type_3d, ierr)

    if (this%send_type_2d /= MPI_DATATYPE_NULL) call MPI_TYPE_FREE(this%send_type_2d, ierr)
    if (this%recv_type_2d /= MPI_DATATYPE_NULL) call MPI_TYPE_FREE(this%recv_type_2d, ierr)

  end subroutine latlon_halo_final

end module latlon_halo_mod
